from dataclasses import dataclass

import numpy as np
import pandas as pd
from pandas import DataFrame

from _index_weight import read_get_data as get_index_weight
from _stock_basic import read_data as read_stock_basic
from app_config import get_engine_ts

pd.set_option('display.max_rows', None)  # 显示所有行
pd.set_option('display.max_columns', None)  # 显示所有列

# 设置显示宽度，防止换行
pd.set_option('display.width', None)  # 自动检测控制台宽度
pd.set_option('display.max_colwidth', None)  # 不截断列内容

engine = get_engine_ts()

_weight_key: str = '权重(%)weight'
column_cons_key: str = "成分券代码Constituent Code"
cs_index_index_name = '指数名称 Index Name'
cs_index_stock_code = '成分券代码Constituent Code'
cs_index_stock_name = '成分券名称Constituent Name'
cs_index_buy_weight = '权重(%)weight'

stock_basic = read_stock_basic()


@dataclass
class FundInfo(object):
    def __init__(self):
        self.name = None
        self.code = None
        self.date = None
        self.index = None
        self.amount = None


class StockInfo(object):
    def __init__(self):
        self.code = None
        self.name = None
        self.amount = None
        self.msg_buy = None
        self.msg_sell = None


def get_simple_ts_code(_param):
    return _param.split('.')[0]
    pass


def fund_amount(fund_info: FundInfo = None) -> dict[str, DataFrame]:
    xls_cons: str = fund_info.code + 'cons.xls'
    next_cons: DataFrame = pd.read_excel("data/" + xls_cons, sheet_name=xls_cons, dtype={column_cons_key: str})
    index_name = next_cons.iloc[0][cs_index_index_name]
    next_cons = index_weight1(fund_info.date, next_cons, column_cons_key)

    cur_cons = get_index_weight(index_code=fund_info.index, date_str=fund_info.date)

    cur_cons[column_cons_key] = cur_cons.apply(lambda row_: get_simple_ts_code(row_['con_code']), axis=1)

    set_buy = set(next_cons[column_cons_key]) - set(cur_cons[column_cons_key])
    set_sell = set(cur_cons[column_cons_key]) - set(next_cons[column_cons_key])

    buy: DataFrame = next_cons[next_cons[column_cons_key].isin(set_buy)].copy()
    sell: DataFrame = cur_cons[cur_cons[column_cons_key].isin(set_sell)].copy()

    buy['buy_amount'] = buy[_weight_key] * fund_info.amount / 100
    sell['sell_amount'] = sell['weight'] * fund_info.amount / 100

    buy = buy.rename(columns={cs_index_index_name: 'index_name',
                              cs_index_stock_code: 'stock_code',
                              cs_index_stock_name: 'stock_name',
                              cs_index_buy_weight: 'buy_weight'})
    buy = buy[['index_name', 'stock_code', 'stock_name', 'buy_weight', 'buy_amount']].copy()

    sell = pd.merge(left=sell, right=stock_basic, left_on='con_code', right_on='ts_code', how='left')
    sell['index_name'] = index_name
    sell['stock_code'] = sell.apply(lambda _row_: get_simple_ts_code(_row_['con_code']), axis=1)
    sell['stock_name'] = sell['name']
    sell['sell_weight'] = sell['weight']
    sell = sell[['index_name', 'stock_code', 'stock_name', 'sell_weight', 'sell_amount']].copy()

    buy.to_excel('bs/' + fund_info.code + '_' + fund_info.name + 'buy.xlsx')
    sell.to_excel('bs/' + fund_info.code + '_' + fund_info.name + 'sell.xlsx')
    sell_ = {"buy": buy, "sell": sell}
    return sell_


def index_weight1(trade_date, df_index, head_key):
    index_con = df_index
    index_con['symbol'] = index_con[head_key]

    query = f""" select symbol, ts_code from stock_basic """
    stock_basic = pd.read_sql(query, engine)

    index_con = pd.merge(index_con, stock_basic, on='symbol', how='left')
    index_con['con_code'] = index_con['ts_code']
    con_codes = index_con['con_code'].tolist()

    # 构建查询语句
    query = f"""
           SELECT ts_code,close,total_share,free_share 
           FROM `daily_basic`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in con_codes)})
           AND trade_date = '{trade_date}'
           """
    # 执行查询并将结果转换为DataFrame
    stock_share = pd.read_sql_query(query, engine)
    index_con_share = pd.merge(index_con, stock_share, left_on="con_code", right_on="ts_code", how='left')

    index_con_share['calc_percent'] = index_con_share['free_share'] / index_con_share['total_share'] * 100

    # 定义生成 b 列的规则
    def generate_b(value):
        if value <= 15:
            return np.ceil(value)  # 上调至最接近的整数
        elif 15 < value <= 20:
            return 20
        elif 20 < value <= 30:
            return 30
        elif 30 < value <= 40:
            return 40
        elif 40 < value <= 50:
            return 50
        elif 50 < value <= 60:
            return 60
        elif 60 < value <= 70:
            return 70
        elif 70 < value <= 80:
            return 80
        else:
            return 100

    index_con_share['adjust_percent'] = index_con_share['calc_percent'].apply(generate_b)
    index_con_share['adjust_share'] = index_con_share['adjust_percent'] / 100 * index_con_share['total_share']
    index_con_share['adjust_amount'] = index_con_share['adjust_share'] * index_con_share['close']

    amount_sum = index_con_share['adjust_amount'].sum()

    index_con_share[_weight_key] = index_con_share['adjust_amount'] / amount_sum * 100
    return index_con_share


def get_index_code(index_basics_: DataFrame) -> str:
    if index_basics_.empty:
        raise ValueError("输入DataFrame不能为空")

    def market_priority(market_value) -> int:
        priority_order = {
            'SSE': 1,
            'SZSE': 2,
            'CSI': 3,
            'CNI': 4,
            'SW': 4,
            'MSCI': 5,
            'CICC': 6,
            'OTH': 7
        }
        return priority_order.get(market_value, 999)

    sorted_df = index_basics_.sort_values(
        'market',
        key=lambda x: x.map(market_priority)
    )
    return sorted_df.iloc[0]['ts_code']


if __name__ == '__main__':
    with open('fund_amount.csv', 'r', encoding='utf-8') as csv_f:
        content = csv_f.readlines()
    excel = pd.read_csv('fund_amount.csv', dtype={'amount': float, 'date': str, 'code': str})

    params: list[FundInfo] = []

    for index, row in excel.iterrows():
        row_: FundInfo = FundInfo()
        row_.name = row['name'].strip()
        row_.amount = row['amount']
        row_.date = row['date'].strip()
        row_.code = row['code'].strip()
        if 'H30184'.__eq__(row_.code) | 'H30590'.__eq__(row_.code):
            continue
        # 构建查询语句
        query = f""" 
        select * from index_basic where ts_code like '{row_.code}.%%' 
        """
        # 执行查询并将结果转换为DataFrame
        index_basics: DataFrame = pd.read_sql_query(query, engine)
        index_code = get_index_code(index_basics)
        row_.index = index_code
        params.append(row_)

    returns: list[dict[str, DataFrame]] = []

    for param in params:
        re = fund_amount(param)
        returns.append(re)

    stock_dict: dict[str, StockInfo] = {}

    for re in returns:
        buy_df: pd.DataFrame = re.get("buy")
        sell_df: pd.DataFrame = re.get("sell")

        for index, row in buy_df.iterrows():
            buy_stock_code = row['stock_code']
            buy_stock_name = row['stock_name']
            buy_index_name = row['index_name']
            buy_weigh = row['buy_weight']
            buy_amount = row['buy_amount']

            sto: StockInfo
            if stock_dict.__contains__(buy_stock_code):
                sto = stock_dict.get(buy_stock_code)
            else:
                sto = StockInfo()
                sto.msg_buy = ''
                sto.msg_sell = ''
                sto.amount = 0.0
                sto.code = buy_stock_code
                sto.name = buy_stock_name
                stock_dict.update({buy_stock_code: sto})

            sto.amount = sto.amount + buy_amount
            # sto.msg_buy = sto.msg_buy + "[" + buy_index_name + ": " + str(buy_amount) + "] "
            sto.msg_buy = f"{sto.msg_buy} [{buy_index_name}: {buy_amount:.1f}]"

        for index, row in sell_df.iterrows():
            buy_stock_code = row['stock_code']
            buy_stock_name = row['stock_name']
            buy_index_name = row['index_name']
            buy_weigh = row['sell_weight']
            buy_amount = row['sell_amount']

            sto: StockInfo
            if stock_dict.__contains__(buy_stock_code):
                sto = stock_dict.get(buy_stock_code)
            else:
                sto = StockInfo()
                sto.msg_buy = ''
                sto.msg_sell = ''
                sto.amount = 0.0
                sto.code = buy_stock_code
                sto.name = buy_stock_name
                stock_dict.update({buy_stock_code: sto})

            sto.amount = sto.amount - buy_amount
            # sto.msg_sell = sto.msg_sell + "[" + buy_index_name + ": " + str(buy_amount) + "] "
            sto.msg_sell = f"{sto.msg_sell} [{buy_index_name}: {buy_amount:.1f}]"

    list_data = [
        {
            "name": p.name,
            "code": p.code,
            "bs_amount": p.amount,
            "buy": p.msg_buy,
            "sell": p.msg_sell
        } for p in stock_dict.values()]
    frame = pd.DataFrame(data=list_data)

    # 构建查询语句2.426248201
    query = f"""
           select sb.symbol as code, sb.ts_code as ts_code from stock_basic sb
           """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    frame = pd.merge(frame, result_df, on='code', how='left')

    filter_ts_code = frame['ts_code'].tolist()
    strStartDate = '20250526'
    strEndDate = '20250530'
    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, amount FROM `daily`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
        AND trade_date >= '{strStartDate}'
        AND trade_date <= '{strEndDate}'
        """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 100000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(frame, grouped_df, on='ts_code', how='left')

    merge['冲击系数'] = merge['bs_amount'] / merge['日均成交额']
    merge = merge.sort_values("冲击系数", ascending=False).reset_index(drop=True)
    merge.to_excel("result3_v2.xlsx")
